package com.kun.ai.common;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.kun.ai.properties.DashScopeProperties;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * @author kun.li
 */
@Component
@RequiredArgsConstructor
public class DashScopeHelper {

    private final DashScopeProperties dashScopeProperties;


    public GenerationResult generationResult(List<Message> messageList) {
        Generation generation = new Generation();
        GenerationParam param = GenerationParam.builder()
                .apiKey(dashScopeProperties.getApiKey())
                .model(dashScopeProperties.getModel())
                .messages(messageList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        GenerationResult result = null;
        try {
            result = generation.call(param);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return result;
    }

    public String getContent(List<Message> messageList) {
        return generationResult(messageList).getOutput().getChoices().get(0).getMessage().getContent();
    }
}
